import time
import os
import matplotlib.pyplot as plt
import numpy as np

# 导入你实现的算法
from publickey.rsa_core import RSA
from publickey.elgamal_core import ElGamal
from publickey.ecc_core import (
    get_curve_by_name,
    generate_keypair as ecc_gen_keys,
    ecc_encrypt_simplified,
    ecc_decrypt_simplified
)

def plot_results(results):
    """将实验结果可视化"""
    # 指定一个支持中文的字体
    plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei'] 
    plt.rcParams['axes.unicode_minus'] = False 

    # --- 1. 密钥生成时间 ---
    labels = list(results['key_gen_time'].keys())
    values = list(results['key_gen_time'].values())
    plt.figure(figsize=(12, 7))
    plt.bar(labels, values, color=['blue', 'green', 'purple', 'red', 'orange', 'brown', 'cyan', 'magenta', 'yellow'])
    plt.ylabel('时间 (秒)')
    plt.title('密钥生成时间比较')
    plt.yscale('log')
    plt.xticks(rotation=45, ha="right")
    plt.tight_layout()
    plt.savefig('key_generation_time.png')
    plt.show()

    # --- 2. 加密/解密时间 (平均值) ---
    labels = list(results['encrypt_time'].keys())
    encrypt_times = list(results['encrypt_time'].values())
    decrypt_times = list(results['decrypt_time'].values())
    
    x = np.arange(len(labels))
    width = 0.35
    
    fig, ax = plt.subplots(figsize=(14, 8))
    rects1 = ax.bar(x - width/2, encrypt_times, width, label='加密', color='skyblue')
    rects2 = ax.bar(x + width/2, decrypt_times, width, label='解密', color='salmon')
    
    ax.set_ylabel('平均时间 (秒)')
    ax.set_title('平均加密与解密时间比较 (100次)')
    ax.set_xticks(x)
    ax.set_xticklabels(labels, rotation=45, ha="right")
    ax.legend()
    ax.set_yscale('log')
    plt.tight_layout()
    plt.savefig('encryption_decryption_time_avg.png')
    plt.show()

def run_experiments():
    """运行所有性能评估实验"""
    
    num_iterations = 100
    print(f"每个算法将进行 {num_iterations} 次加密/解密操作来计算平均时间。")
    
    # --- 参数设置 ---
    rsa_elgamal_key_sizes = [1024, 2048, 3072] 
    ecc_curve_names = ["secp256k1", "secp256r1", "secp384r1"]
    message_size_bytes = 40

    results = {
        "key_gen_time": {},
        "encrypt_time": {},
        "decrypt_time": {},
        "ciphertext_size": {}
    }

    # --- 1. RSA 实验 ---
    print("\n" + "="*10, "开始 RSA 实验", "="*10)
    messages_rsa = [os.urandom(message_size_bytes) for _ in range(num_iterations)]
    for bits in rsa_elgamal_key_sizes:
        print(f"\n--- 测试 RSA with {bits}-bit key ---")
        algo_id = f"RSA-{bits}"
        start_time = time.time()
        rsa_instance = RSA(bits=bits)
        results["key_gen_time"][algo_id] = time.time() - start_time
        print(f"密钥生成时间: {results['key_gen_time'][algo_id]:.6f} 秒")
        start_time = time.time()
        ciphertexts_rsa = [rsa_instance.encrypt(m) for m in messages_rsa]
        results["encrypt_time"][algo_id] = (time.time() - start_time) / num_iterations
        print(f"加密平均时间: {results['encrypt_time'][algo_id]:.8f} 秒")
        start_time = time.time()
        for c in ciphertexts_rsa: rsa_instance.decrypt(c)
        results["decrypt_time"][algo_id] = (time.time() - start_time) / num_iterations
        print(f"解密平均时间: {results['decrypt_time'][algo_id]:.8f} 秒")
        n, _ = rsa_instance.public_key
        results["ciphertext_size"][algo_id] = (message_size_bytes, (n.bit_length() + 7) // 8)

    # --- 2. ElGamal 实验 ---
    print("\n" + "="*10, "开始 ElGamal 实验", "="*10)
    messages_elgamal_int = [int.from_bytes(os.urandom(message_size_bytes), 'big') for _ in range(num_iterations)]
    for bits in rsa_elgamal_key_sizes:
        # 修改：移除了对 key_size > 2048 的跳过逻辑，以测试3072位ElGamal
        print(f"\n--- 测试 ElGamal with {bits}-bit key ---")
        algo_id = f"ElGamal-{bits}"

        start_time = time.time()
        elgamal_instance = ElGamal(key_size_bits=bits)
        results["key_gen_time"][algo_id] = time.time() - start_time
        print(f"密钥生成时间: {results['key_gen_time'][algo_id]:.6f} 秒")
        
        pub_key = elgamal_instance.public_key
        priv_key = elgamal_instance.private_key
        public_params = (pub_key[0], pub_key[1], pub_key[2])
        
        start_time = time.time()
        ciphertexts_elgamal = [ElGamal.encrypt(pub_key, m) for m in messages_elgamal_int]
        results["encrypt_time"][algo_id] = (time.time() - start_time) / num_iterations
        print(f"加密平均时间: {results['encrypt_time'][algo_id]:.8f} 秒")
        
        start_time = time.time()
        for c in ciphertexts_elgamal: ElGamal.decrypt(priv_key, public_params, c)
        results["decrypt_time"][algo_id] = (time.time() - start_time) / num_iterations
        print(f"解密平均时间: {results['decrypt_time'][algo_id]:.8f} 秒")
        
        p, _, _, _ = pub_key
        results["ciphertext_size"][algo_id] = (message_size_bytes, 2 * ((p.bit_length() + 7) // 8))

    # --- 3. ECC 实验 ---
    print("\n" + "="*10, "开始 ECC 实验", "="*10)
    messages_ecc = [os.urandom(message_size_bytes) for _ in range(num_iterations)]
    for curve_name in ecc_curve_names:
        print(f"\n--- 测试 ECC with {curve_name} ---")
        algo_id = f"ECC-{curve_name}"

        try:
            curve = get_curve_by_name(curve_name)
        except ValueError as e:
            print(e)
            continue

        start_time = time.time()
        ecc_priv_key, ecc_pub_key = ecc_gen_keys(curve)
        results["key_gen_time"][algo_id] = time.time() - start_time
        print(f"密钥生成时间: {results['key_gen_time'][algo_id]:.6f} 秒")

        start_time = time.time()
        ciphertexts_ecc = [ecc_encrypt_simplified(ecc_pub_key, curve, m) for m in messages_ecc]
        results["encrypt_time"][algo_id] = (time.time() - start_time) / num_iterations
        print(f"加密平均时间: {results['encrypt_time'][algo_id]:.8f} 秒")

        start_time = time.time()
        for R, C in ciphertexts_ecc: ecc_decrypt_simplified(ecc_priv_key, curve, R, C)
        results["decrypt_time"][algo_id] = (time.time() - start_time) / num_iterations
        print(f"解密平均时间: {results['decrypt_time'][algo_id]:.8f} 秒")

        point_size_bytes = 1 + 2 * ((curve.p.bit_length() + 7) // 8)
        cipher_len = point_size_bytes + message_size_bytes 
        results["ciphertext_size"][algo_id] = (message_size_bytes, cipher_len)
    
    return results

if __name__ == "__main__":
    experiment_results = run_experiments()
    # 打印最终结果摘要
    print("\n\n" + "="*20, "实验结果摘要", "="*20)
    print("\n密钥生成时间 (秒):")
    for k, v in experiment_results["key_gen_time"].items():
        print(f"  {k}: {v:.6f}")
    
    print("\n平均加密时间 (秒):")
    for k, v in experiment_results["encrypt_time"].items():
        print(f"  {k}: {v:.8f}")

    print("\n平均解密时间 (秒):")
    for k, v in experiment_results["decrypt_time"].items():
        print(f"  {k}: {v:.8f}")
        
    print("\n数据扩展性 (明文 -> 密文, 字节):")
    for k, v in experiment_results["ciphertext_size"].items():
        print(f"  {k}: {v[0]} -> {v[1]}")
        
    plot_results(experiment_results)